package controller

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/go-logr/logr"
	batchv1 "k8s.io/api/batch/v1"
	k8sapierror "k8s.io/apimachinery/pkg/api/errors"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/builder"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"

	"github.com/aquasecurity/trivy-operator/pkg/apis/aquasecurity/v1alpha1"
	"github.com/aquasecurity/trivy-operator/pkg/docker"
	"github.com/aquasecurity/trivy-operator/pkg/exposedsecretreport"
	"github.com/aquasecurity/trivy-operator/pkg/kube"
	"github.com/aquasecurity/trivy-operator/pkg/operator/etc"
	"github.com/aquasecurity/trivy-operator/pkg/operator/jobs"
	"github.com/aquasecurity/trivy-operator/pkg/operator/workload"
	"github.com/aquasecurity/trivy-operator/pkg/plugins/trivy"
	"github.com/aquasecurity/trivy-operator/pkg/sbomreport"
	"github.com/aquasecurity/trivy-operator/pkg/trivyoperator"
	"github.com/aquasecurity/trivy-operator/pkg/vulnerabilityreport"

	. "github.com/aquasecurity/trivy-operator/pkg/operator/predicate"
)

const trivyServerUp = "trivy_server_up"

// WorkloadController watches Kubernetes workloads and generates
// v1alpha1.VulnerabilityReport instances using vulnerability scanner that that
// implements the Plugin interface.
type WorkloadController struct {
	logr.Logger
	etc.Config
	kube.ObjectResolver
	client.Client
	jobs.LimitChecker
	vulnerabilityreport.Plugin
	trivyoperator.PluginContext
	kube.SecretsReader
	trivyoperator.ConfigData
	ServerHealthChecker     ServerHealthChecker
	VulnerabilityReadWriter vulnerabilityreport.ReadWriter
	ExposedSecretReadWriter exposedsecretreport.ReadWriter
	SbomReadWriter          sbomreport.ReadWriter
	SubmitScanJobChan       chan ScanJobRequest
	ResultScanJobChan       chan ScanJobResult
	CacheSyncTimeout        time.Duration
}

// ScanJobResult encapsulate processing result and error
type ScanJobResult struct {
	Result ctrl.Result
	Error  error
}

// ScanJobRequest encapsulate workload and context for processing
type ScanJobRequest struct {
	Workload          client.Object
	Context           context.Context
	ClusterSbomReport map[string]v1alpha1.SbomReportData
}

// +kubebuilder:rbac:groups="",resources=pods,verbs=get;list;watch
// +kubebuilder:rbac:groups=apps,resources=replicasets,verbs=get;list;watch
// +kubebuilder:rbac:groups="",resources=replicationcontrollers,verbs=get;list;watch
// +kubebuilder:rbac:groups=apps,resources=statefulsets,verbs=get;list;watch
// +kubebuilder:rbac:groups=apps,resources=daemonsets,verbs=get;list;watch
// +kubebuilder:rbac:groups=batch,resources=cronjobs,verbs=get;list;watch
// +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch
// +kubebuilder:rbac:groups=aquasecurity.github.io,resources=vulnerabilityreports,verbs=get;list;watch;create;update;patch;delete
// +kubebuilder:rbac:groups=aquasecurity.github.io,resources=exposedsecretreports,verbs=get;list;watch;create;update;patch;delete
// +kubebuilder:rbac:groups=aquasecurity.github.io,resources=sbomreports,verbs=get;list;watch;create;update;patch;delete
// +kubebuilder:rbac:groups=aquasecurity.github.io,resources=clustersbomreports,verbs=get;list;watch;create;update;patch;delete

// Manage scan jobs with image pull secrets
// kubebuilder:rbac:groups="",resources=secrets,verbs=create;update

func (r *WorkloadController) SetupWithManager(mgr ctrl.Manager) error {
	installModePredicate, err := InstallModePredicate(r.Config)
	if err != nil {
		return err
	}

	var resources []kube.Resource

	// Determine which Kubernetes workloads the controller will reconcile and add them to resources
	targetWorkloads := r.Config.GetTargetWorkloads()
	for _, tw := range targetWorkloads {
		var resource kube.Resource
		err := resource.GetWorkloadResource(tw, nil, r.ObjectResolver)
		if err != nil {
			return err
		}
		resources = append(resources, resource)
	}

	for _, resource := range resources {
		err = ctrl.NewControllerManagedBy(mgr).WithOptions(controller.Options{
			CacheSyncTimeout: r.CacheSyncTimeout,
		}).
			For(resource.ForObject, builder.WithPredicates(
				Not(ManagedByTrivyOperator),
				Not(IsBeingTerminated),
				installModePredicate,
			)).
			Owns(&v1alpha1.VulnerabilityReport{}).
			Owns(&v1alpha1.ExposedSecretReport{}).
			Owns(&v1alpha1.SbomReport{}).
			Complete(r.reconcileWorkload(resource.Kind))
		if err != nil {
			return err
		}
	}
	// process scan jobs
	go r.ProcessScanJob()
	return nil
}

func (r *WorkloadController) reconcileWorkload(workloadKind kube.Kind) reconcile.Func {
	return func(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
		log := r.Logger.WithValues("kind", workloadKind, "name", req.NamespacedName)
		workloadRef := kube.ObjectRefFromKindAndObjectKey(workloadKind, req.NamespacedName)
		log.V(1).Info("Getting workload from cache", "workload", workloadRef.Name)
		workloadObj, err := r.ObjectFromObjectRef(ctx, workloadRef)
		if err != nil {
			if k8sapierror.IsNotFound(err) {
				log.V(1).Info("Ignoring cached workload that must have been deleted", "workload", workloadRef.Name)
				return ctrl.Result{}, nil
			}
			return ctrl.Result{}, fmt.Errorf("getting %s from cache: %w", workloadKind, err)
		}

		// Skip processing if it's a Pod controlled by a built-in K8s workload.
		if skip, err := workload.SkipProcessing(ctx, workloadObj, r.ObjectResolver,
			r.Config.VulnerabilityScannerScanOnlyCurrentRevisions, log, r.ConfigData.GetSkipResourceByLabels()); skip {
			return ctrl.Result{}, err
		}

		podSpec, err := kube.GetPodSpec(workloadObj)
		if err != nil {
			return ctrl.Result{}, err
		}

		containerImages := kube.GetContainerImagesFromPodSpec(podSpec, r.GetSkipInitContainers())
		hash := kube.ComputeHash(podSpec)

		hasVulnReports := true
		if r.Config.VulnerabilityScannerEnabled {
			hasVulnReports, err = hasVulnerabilityReports(ctx, r.VulnerabilityReadWriter, workloadRef, hash, containerImages)
			if err != nil {
				return ctrl.Result{}, err
			}
		}
		hasExposedSecretReports := true
		if r.Config.ExposedSecretScannerEnabled {
			hasExposedSecretReports, err = hasSecretReports(ctx, r.ExposedSecretReadWriter, workloadRef, hash, containerImages)
			if err != nil {
				return ctrl.Result{}, err
			}
		}

		if hasVulnReports && hasExposedSecretReports {
			log.V(1).Info("VulnerabilityReports or Secret Reports already exist for the workload", "workload", workloadRef.Name)
			return ctrl.Result{}, nil
		}
		var reportsData map[string]v1alpha1.SbomReportData

		_, job, err := r.hasActiveScanJob(ctx, workloadRef, hash)
		if err != nil {
			return ctrl.Result{}, fmt.Errorf("checking scan job: %w", err)
		}

		if job != nil {
			log.V(1).Info("Scan job already exists",
				"job", fmt.Sprintf("%s/%s", job.Namespace, job.Name))
			return ctrl.Result{}, nil
		}

		if r.BuiltInTrivyServer {
			trivyServerAvaliable, err := r.ServerHealthChecker.TrivyServerAvaliable(r.GeTrivyServerURL())
			if err != nil {
				return ctrl.Result{}, err
			}
			if !trivyServerAvaliable {
				return ctrl.Result{RequeueAfter: r.Config.ScanJobRetryAfter}, nil
			}
		}
		if r.Config.ClusterSbomCacheEnable {
			reportsData = getGlobalSbomReports(ctx, r.SbomReadWriter, containerImages, log)
			if len(reportsData) > 0 {
				err = r.reuseSbomReport(ctx, workloadObj, reportsData)
				if err != nil {
					log.V(1).Info("unable to reuse report for image containers from cached sbom  - running standard scan-job", "image containers", containerImages)
					reportsData = make(map[string]v1alpha1.SbomReportData)
				}
			}
		}
		log.V(1).Info("Submitting a scan for the workload", "workload", workloadRef.Name)
		// sync all potential workload for scanning
		r.SubmitScanJobChan <- ScanJobRequest{Workload: workloadObj, Context: ctx, ClusterSbomReport: reportsData}
		// collect scan job processing results
		scanJobResult := <-r.ResultScanJobChan
		return scanJobResult.Result, scanJobResult.Error
	}
}

func (r *WorkloadController) ProcessScanJob() {
	for workloadRequest := range r.SubmitScanJobChan {
		log := r.Logger.WithValues("pre scan job processing for workload:", workloadRequest.Workload.GetName())
		limitExceeded, scanJobsCount, err := r.LimitChecker.Check(workloadRequest.Context)
		if err != nil {
			r.ResultScanJobChan <- ScanJobResult{Result: ctrl.Result{}, Error: err}
			continue
		}
		log.V(1).Info("Checking scan jobs limit", "count", scanJobsCount, "limit", r.ConcurrentScanJobsLimit)

		if limitExceeded {
			log.V(1).Info("Pushing back scan job", "count", scanJobsCount, "retryAfter", r.ScanJobRetryAfter)
			r.ResultScanJobChan <- ScanJobResult{Result: ctrl.Result{RequeueAfter: r.Config.ScanJobRetryAfter}, Error: nil}
			continue
		}
		err = r.SubmitScanJob(workloadRequest.Context, workloadRequest.Workload, workloadRequest.ClusterSbomReport)
		r.ResultScanJobChan <- ScanJobResult{Result: ctrl.Result{}, Error: err}
	}
}

func (r *WorkloadController) hasActiveScanJob(ctx context.Context, owner kube.ObjectRef, hash string) (bool, *batchv1.Job, error) {
	jobName := fmt.Sprintf("scan-vulnerabilityreport-%s", kube.ComputeHash(owner))
	job := &batchv1.Job{}
	err := r.Get(ctx, client.ObjectKey{Namespace: r.Config.Namespace, Name: jobName}, job)
	if err != nil {
		if k8sapierror.IsNotFound(err) {
			return false, nil, nil
		}
		return false, nil, fmt.Errorf("getting job from cache: %w", err)
	}
	if job.Labels[trivyoperator.LabelResourceSpecHash] == hash {
		return true, job, nil
	}
	return false, nil, nil
}

func (r *WorkloadController) SubmitScanJob(ctx context.Context, owner client.Object, reusedReports map[string]v1alpha1.SbomReportData) error {

	log := r.Logger.WithValues("kind", owner.GetObjectKind().GroupVersionKind().Kind,
		"name", owner.GetName(), "namespace", owner.GetNamespace())
	var err error

	credentials := make(map[string]docker.Auth, 0)
	if len(reusedReports) == 0 {
		privateRegistrySecrets, err := r.Config.GetPrivateRegistryScanSecretsNames()
		if err != nil {
			return err
		}

		pConfig, err := r.PluginContext.GetConfig()
		if err != nil {
			return err
		}

		multiSecretSupport := trivy.MultiSecretSupport(trivy.Config{PluginConfig: pConfig})

		if r.AccessGlobalSecretsAndServiceAccount {
			// Global access is enabled - therefore imagePullSecrets from the podSpec can be used
			credentials, err = r.CredentialsByServer(ctx, owner, privateRegistrySecrets, multiSecretSupport, true)
			if err != nil {
				return err
			}
		} else {
			// Global access is disabled - check if privateRegistrySecrets references an imagePullSecret in the operator namespace
			imagePullSecretName, ok := privateRegistrySecrets[r.Config.Namespace]
			if ok {
				// privateRegistrySecrets references an imagePullSecret in the operator namespace - use this for pulling private images
				credentials, err = r.CredentialsByServer(ctx, owner, map[string]string{r.Config.Namespace: imagePullSecretName}, multiSecretSupport, false)
				if err != nil {
					return err
				}
			}
		}
	}

	scanJobAffinity, err := r.GetScanJobAffinity()
	if err != nil {
		return fmt.Errorf("getting scan job affinity: %w", err)
	}

	scanJobTolerations, err := r.GetScanJobTolerations()
	if err != nil {
		return fmt.Errorf("getting scan job tolerations: %w", err)
	}

	scanJobAnnotations, err := r.GetScanJobAnnotations()
	if err != nil {
		return fmt.Errorf("getting scan job annotations: %w", err)
	}

	scanJobNodeSelector, err := r.GetScanJobNodeSelector()
	if err != nil {
		return fmt.Errorf("getting scan job nodeSelector: %w", err)
	}

	scanJobSecurityContext, err := r.GetScanJobPodSecurityContext()
	if err != nil {
		return fmt.Errorf("getting scan job podSecurityContext: %w", err)
	}

	scanJobContainerSecurityContext, err := r.GetScanJobContainerSecurityContext()
	if err != nil {
		return fmt.Errorf("getting scan job [container] securityContext: %w", err)
	}

	scanJobPodTemplateLabels, err := r.GetScanJobPodTemplateLabels()
	if err != nil {
		return fmt.Errorf("getting scan job template labels: %w", err)
	}

	scanJobPodPriorityClassName, err := r.GetScanJobPodPriorityClassName()
	if err != nil {
		return fmt.Errorf("getting scan job priority class name: %w", err)
	}

	cvm, err := r.GetScanJobCustomVolumeMounts()
	if err != nil {
		return fmt.Errorf("getting scan job scan job custom volume mounts: %w", err)
	}
	cv, err := r.GetScanJobCustomVolumes()
	if err != nil {
		return fmt.Errorf("getting scan job scan job custom volume: %w", err)
	}

	scanJobBuilder := vulnerabilityreport.NewScanJobBuilder().
		WithPlugin(r.Plugin).
		WithPluginContext(r.PluginContext).
		WithTimeout(r.Config.ScanJobTimeout).
		WithTTL(r.Config.ScanJobTTL).
		WithScanSecretTTL(r.Config.ScanSecretTTL).
		WithObject(owner).
		WithAffinity(scanJobAffinity).
		WithTolerations(scanJobTolerations).
		WithAnnotations(scanJobAnnotations).
		WithNodeSelector(scanJobNodeSelector).
		WithPodSecurityContext(scanJobSecurityContext).
		WithCustomVolumes(cv).
		WithCustomVolumesMount(cvm).
		WithSecurityContext(scanJobContainerSecurityContext).
		WithSkipInitContainers(r.GetSkipInitContainers()).
		WithPodTemplateLabels(scanJobPodTemplateLabels).
		WithCredentials(credentials).
		WithPodPriorityClassName(scanJobPodPriorityClassName)
	if len(reusedReports) > 0 {
		scanJobBuilder.WithSbomClusterReports(reusedReports)
	}
	scanJob, secrets, err := scanJobBuilder.Get()

	if err != nil {
		if errors.Is(err, kube.ErrReplicaSetNotFound) || errors.Is(err, kube.ErrNoRunningPods) ||
			errors.Is(err, kube.ErrUnSupportedKind) {
			log.V(1).Info("ignoring vulnerability scan", "reason", err)
			return nil
		}
		return fmt.Errorf("constructing scan job: %w", err)
	}
	// If there are no containers in the pod spec, we don't need to create a scan job.
	if len(scanJob.Spec.Template.Spec.Containers) == 0 {
		log.V(1).Info("ignoring vulnerability scan", "no containers in the pod spec for resource", owner.GetName())
		return nil
	}

	for _, secret := range secrets {
		err = r.Client.Create(ctx, secret)
		if err != nil {
			if k8sapierror.IsAlreadyExists(err) {
				return nil
			}
			return fmt.Errorf("creating secret used by scan job failed: %s: %w", secret.Namespace+"/"+secret.Name, err)
		}
	}

	log = log.WithValues("podSpecHash", scanJob.Labels[trivyoperator.LabelResourceSpecHash])

	log.V(1).Info("Creating scan job for the workload")
	err = r.Client.Create(ctx, scanJob)
	if err != nil {
		if k8sapierror.IsAlreadyExists(err) {
			// TODO Delete secrets that were created in the previous step. Alternatively we can delete them on schedule.
			return nil
		}
		return fmt.Errorf("creating scan job failed: %s: %w", scanJob.Namespace+"/"+scanJob.Name, err)
	}

	for _, secret := range secrets {
		err = controllerutil.SetOwnerReference(scanJob, secret, r.Client.Scheme())
		if err != nil {
			return fmt.Errorf("setting owner reference: %w", err)
		}
		err := r.Client.Update(ctx, secret)
		if err != nil {
			return fmt.Errorf("setting owner reference of secret used by scan job failed: %s: %w", secret.Namespace+"/"+secret.Name, err)
		}
	}

	return nil
}

func (r *WorkloadController) reuseSbomReport(ctx context.Context, owner client.Object, sbomReportDataMap map[string]v1alpha1.SbomReportData) error {
	hash, err := kube.ComputeSpecHash(owner)
	if err != nil {
		return err
	}
	resourceLabelsToInclude := r.GetReportResourceLabels()
	additionalCustomLabels, err := r.GetAdditionalReportLabels()
	if err != nil {
		return err
	}
	sbomReports := make([]v1alpha1.SbomReport, 0)
	for containerName, sbomReportData := range sbomReportDataMap {
		sbomReportBuilder := sbomreport.NewReportBuilder(r.Client.Scheme()).
			Controller(owner).
			Container(containerName).
			Data(sbomReportData).
			PodSpecHash(hash).
			ResourceLabelsToInclude(resourceLabelsToInclude).
			AdditionalReportLabels(additionalCustomLabels)
		sbomReport, err := sbomReportBuilder.NamespacedReport()
		if err != nil {
			return err
		}
		sbomReports = append(sbomReports, sbomReport)
	}
	return r.SbomReadWriter.Write(ctx, sbomReports)
}
